package org.hepeng.workx.spring.session.cookie;

import com.google.common.net.InternetDomainName;
import lombok.Data;
import org.apache.commons.lang3.StringUtils;
import org.joor.Reflect;
import org.springframework.session.web.http.DefaultCookieSerializer;
import sun.net.util.IPAddressUtil;

import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * @author he peng
 */

@Data
public class TopLevelDomainCookieSerializer extends DefaultCookieSerializer {

    private static final Pattern DEFAULT_TOP_LEVEL_DOMAIN_PATTERN = Pattern.compile("(?<=http://|\\.)[^.]*?\\.(com|cn|net|org|biz|info|cc|tv)",Pattern.CASE_INSENSITIVE);

    private Pattern topLevelDomainPattern = DEFAULT_TOP_LEVEL_DOMAIN_PATTERN;

    @Override
    public void writeCookieValue(CookieValue cookieValue) {
        HttpServletRequest request = cookieValue.getRequest();
        HttpServletResponse response = cookieValue.getResponse();

        String jvmRoute = Reflect.on(this).get("jvmRoute");
        String requestedCookieValue = cookieValue.getCookieValue();
        String actualCookieValue = jvmRoute == null ? requestedCookieValue
                : requestedCookieValue + jvmRoute;

        String cookieName = Reflect.on(this).get("cookieName");
        boolean useBase64Encoding = Reflect.on(this).get("useBase64Encoding");

        String base64EncodeCookieValue = Reflect.on(this).call("base64Encode" , actualCookieValue).get();
        Cookie sessionCookie = new Cookie(cookieName , useBase64Encoding
                ? base64EncodeCookieValue : actualCookieValue);

        boolean secure = Reflect.on(this).call("isSecureCookie" , request).get();
        String cookiePath = Reflect.on(this).call("getCookiePath" , request).get();
        sessionCookie.setSecure(secure);
        sessionCookie.setPath(cookiePath);
        String domainName = getTopLevelDomainName(request);
        if (domainName != null) {
            sessionCookie.setDomain(domainName);
        }

        boolean useHttpOnlyCookie = Reflect.on(this).get("useHttpOnlyCookie");
        if (useHttpOnlyCookie) {
            sessionCookie.setHttpOnly(true);
        }

        String rememberMeRequestAttribute = Reflect.on(this).get("rememberMeRequestAttribute");
        Integer cookieMaxAge = Reflect.on(this).get("cookieMaxAge");
        if ("".equals(requestedCookieValue)) {
            sessionCookie.setMaxAge(0);
        } else if (rememberMeRequestAttribute != null
                && request.getAttribute(rememberMeRequestAttribute) != null) {
            // the cookie is only written at time of session creation, so we rely on
            // session expiration rather than cookie expiration if remember me is enabled
            sessionCookie.setMaxAge(Integer.MAX_VALUE);
        } else {
            sessionCookie.setMaxAge(cookieMaxAge);
        }

        response.addCookie(sessionCookie);
    }

    private String getTopLevelDomainName(HttpServletRequest request) {
        String domainName = Reflect.on(this).call("getDomainName" , request).get();
        if (StringUtils.isBlank(domainName)) {
            domainName = request.getServerName();
        }
        if (IPAddressUtil.isIPv4LiteralAddress(domainName)) {
            return domainName;
        }

        if (! InternetDomainName.isValid(domainName)) {
            return domainName;
        }

        InternetDomainName internetDomainName = InternetDomainName.from(domainName);
        return internetDomainName.topPrivateDomain().toString();
    }
}
